"""SPH Interpolation kernels"""

import numpy

SPLINEFAC1  = 2.0/3.0

def cs_function(xa, xb, h):
    """Cubic B-Spline kernel.
    """
    rab = abs( xa - xb )
    h1 = 1.0/h
    q = rab * h1
    val = 0.0

    if q > 2.0:
        val = 0.0
    elif q > 1.0:
        val = 0.25 * ( 2-q ) * ( 2-q ) * ( 2-q )
    else:
        val = 1 - 1.5 * q * q * ( 1 - 0.5 * q )

    return SPLINEFAC1 * h1 * val

class SPHInterpolate(object):

    @classmethod
    def get_density(self, nnps, pnts, h):

        n = len(pnts)
        rho = numpy.zeros(pnts.shape)
        
        for i in range(n):
            xa = pnts[i]

            # get neighbors for this particle
            neighbors = nnps.get_nearest_particles( xa )

            for j in neighbors:
                xb = pnts[j]
                rho[i] += cs_function(xa, xb, h)

        return rho

    @classmethod
    def interpolate(self, nnps, dst, src, data, rho, h):

        result = numpy.zeros(dst.shape)

        nd = len(dst)
        for i in range(nd):
            xa = dst[i]

            # get neighbors for this particle
            neighbors = nnps.get_nearest_particles( xa )

            for j in neighbors:
                xb = src[j]; fb = data[j]; rb = rho[j]
                result[i] += fb/rb * cs_function(xa, xb, h)            

        return result

class NNPS:

    def bin(self, grid, cell_size=None):
        """Bin the particles"""
        self.cell_size = cell_size

        np = grid.ncells
        x = grid.xc

        # reset the bin structure
        self.cells = cells = {}

        #now bin the particles
        for i in range(np):
            _x = x[i]
            cid = int(_x/cell_size)

            if not cells.has_key( cid ):
                cells[cid] = []

            indices = cells[cid]
            indices.append( i )

    def get_nearest_particles(self, xa):
        """Return neighbor indices for the specified particle"""

        # get the bin structure
        cells = self.cells

        # find the cid for the requested particle
        _cid = int( xa/self.cell_size )

        neighbors = []

        # append the neighbors
        for cid in [_cid - 1, _cid, _cid + 1]:

            indices = cells.get( cid )
            if indices is not None:
                neighbors.extend( indices )

        return neighbors
